Kth Smallest Element in a BST

Given a binary search tree, write a function kthSmallest to find the kth smallest element in it.

Note:

You may assume k is always valid, 1 ≤ k ≤ BST’s total elements.

Follow up:

What if the BST is modified (insert/delete operations) often and you need to find the kth smallest frequently? How would you optimize the kthSmallest routine?

Hint:

  1. Try to utilize the property of a BST.
  2. What if you could modify the BST node’s structure?
  3. The optimal runtime complexity is O(height of BST).

Solution:

  1. /**
  2.  * Definition for a binary tree node.
  3.  * public class TreeNode {
  4.  *     int val;
  5.  *     TreeNode left;
  6.  *     TreeNode right;
  7.  *     TreeNode(int x) { val = x; }
  8.  * }
  9.  */
  10. public class Solution {
  11.   public int kthSmallest(TreeNode root, int k) {
  12.     int nLeft = countNodes(root.left);
  14.     if (nLeft == k - 1)
  15.       return root.val;
  17.     else if (nLeft > k - 1)
  18.       return kthSmallest(root.left, k);
  20.     else
  21.       return kthSmallest(root.right, k - nLeft - 1);
  22.   }
  24.   int countNodes(TreeNode root) {
  25.     if (root == null)
  26.       return 0;
  28.     return 1 + countNodes(root.left) + countNodes(root.right);
  29.   }
  30. }